import os.path as osp
from PIL import Image
from numpy.core.machar import MachAr
from torch.utils.data import Dataset
from torchvision import transforms
import os
import numpy as np
import json
import torch
from torchvision.transforms.transforms import ToTensor

class Recognition36(Dataset):

    def __init__(self, setname, args):
        IMAGE_PATH = os.path.join(args.data_dir, 'recognition36/images')
        SPLIT_PATH = os.path.join(args.data_dir, 'recognition36')
        if setname in ["test", "novel_all"]:
            json_path = osp.join(SPLIT_PATH, 'novel_all.json') # 当测试数据集与源数据集不一致时使用
        else:
            json_path = osp.join(SPLIT_PATH, setname + '.json')
        with open(json_path, "r") as f:
            self.meta = json.load(f) # json file(dict) ： {"label_names:[...], "image_names":[...], "image_labels":[...]}

        if args.use_specific_status:
            # recogniton36 数据集均值方差
            self.MEAN = np.array((0.43831334, 0.43955493, 0.44021428))
            self.STD = np.array((0.2708545, 0.2706524, 0.27064976))
        else:
            # imagenet 均值方差
            self.MEAN = np.array([x / 255.0 for x in [125.3, 123.0, 113.9]])
            self.STD = np.array([x / 255.0 for x in [63.0, 62.1, 66.7]])


        data = self.meta["image_names"]
        label = self.meta["image_labels"] # dataset 返回的label并无实际作用,只是当做一个类别标识
        data = [_.replace("filelists", "datasets") for _ in data]
 
        self.data = data  # data path of all data
        self.label = label  # label of all data
        self.num_class = len(set(label))

        if setname == 'val' or setname == 'test':
            image_size = args.image_size
            self.transform = transforms.Compose([
                transforms.Resize(int(image_size / 0.9)),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                transforms.Normalize(self.MEAN, self.STD)
                                    ])
        elif setname == 'train':
            image_size = args.image_size
            self.transform = transforms.Compose([
                transforms.RandomResizedCrop(image_size),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(self.MEAN, self.STD)
                                     ])
        try:
            if args.mixmatch:
                # 由于 unlabel 样本需要Ｋ次增强，因此覆盖增强方法，在数据集外部进行增强, 只保留 Resize 方法
                self.transform = transforms.Compose(
                    [
                        transforms.Resize(int(image_size / 0.9)),
                        transforms.CenterCrop(image_size),
                        transforms.ToTensor(),
                        transforms.Normalize(self.MEAN, self.STD)
                    ]
                )
        except:
            pass

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        # 用于处理当 unlabel + query > 类别数目时的情况
        if i == -1:
            return torch.zeros([3, self.img_size, self.img_size]), 0
        path, label = self.data[i], self.label[i]
        image = self.transform(Image.open(path).convert('RGB'))
        return image, 1


if __name__ == '__main__':
    pass